from langchain_chroma import Chroma
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnableLambda, RunnablePassthrough
from langchain_ollama import OllamaEmbeddings
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_community.chat_message_histories import ChatMessageHistory

model = ChatOpenAI(model="ictrek/qwen7b:32k",
                   openai_api_key="ollama",
                   openai_api_base="http://10.2.4.31:11434/v1/")

# 准备数据
docs = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好闻名",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="猫是独立的宠物，通常喜欢自己的空间",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼是初学者的流行宠物，需要相对简单的护理",
        metadata={"source": "鱼类宠物文档"}
    )
]

# 实例化向量数据空间
# pip install -qU langchain-ollama
# 使用本地的嵌入模型
# openAI默认的嵌入模型：OpenAIEmbeddings()
embeddings = OllamaEmbeddings(
    model="shaw/dmeta-embedding-zh",
    base_url="http://10.2.4.31:11434",
)
vector_store = Chroma.from_documents(documents=docs, embedding=embeddings)

# 相似度查询，包括相似度分值，分数越低，相似度越高
# res = vector_store.similarity_search_with_score("咖啡猫")
# print(res)

# 得到检索器对象： bind(k=1) 返回相似度最高的数据
query = RunnableLambda(vector_store.similarity_search).bind(k=1)

# print(query.batch(["咖啡猫"]))  # 检索

# 提示模版
messages = """
使用提供的上下文仅回答这个问题。
{question}
上下文：
{context}
"""

prompt_template = ChatPromptTemplate.from_messages([
    ("human", messages)
])

# RunnablePassthrough: 允许将用户的问题之后传递给prompt和model
chain = {"question": RunnablePassthrough(), "context": query} | prompt_template | model

resp = chain.invoke('请介绍一下猫')

print(resp.content)

